#!/bin/bash
set -e
trap 'echo -e "\n${RED}已中断，脚本退出${NC}"; exit 1' INT

# Trivy 管理脚本 - 固定缓存与结果目录版

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
NC='\033[0m' # 恢复默认颜色

# 固定目录定义（核心修改：指定缓存和结果目录）
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
FIXED_CACHE_DIR="$SCRIPT_DIR/cache"
FIXED_RESULT_DIR="$SCRIPT_DIR/result"

# 检查命令是否存在
check_command() {
    if ! command -v "$1" &> /dev/null; then
        echo -e "${RED}[✗] 错误: $1 未安装！${NC}" >&2
        exit 1
    fi
}

# 获取固定缓存目录（仅返回纯路径）
get_cache_dir() {
    echo -e "${BLUE}[*] 强制使用缓存目录: $FIXED_CACHE_DIR${NC}" >&2

    # 确保目录存在
    if [ ! -d "$FIXED_CACHE_DIR" ]; then
        echo -e "${YELLOW}[!] 缓存目录不存在，尝试创建...${NC}" >&2
        mkdir -p "$FIXED_CACHE_DIR"
        chmod -R 777 "$FIXED_CACHE_DIR"
    fi

    # 检查目录是否可写
    if [ -w "$FIXED_CACHE_DIR" ]; then
        echo "$FIXED_CACHE_DIR"
        return 0
    else
        echo -e "${RED}[✗] 缓存目录 $FIXED_CACHE_DIR 不可写！${NC}" >&2
        echo -e "${YELLOW}[!] 尝试修复权限...${NC}" >&2
        chmod -R 777 "$FIXED_CACHE_DIR"
        
        if [ -w "$FIXED_CACHE_DIR" ]; then
            echo -e "${GREEN}[✓] 权限修复成功${NC}" >&2
            echo "$FIXED_CACHE_DIR"
            return 0
        else
            echo -e "${RED}[✗] 权限修复失败，请手动执行: sudo chmod -R 777 $FIXED_CACHE_DIR${NC}" >&2
            exit 1
        fi
    fi
}

# 创建结果导出目录（固定路径）
create_results_dir() {
    echo -e "${BLUE}[*] 强制使用结果目录: $FIXED_RESULT_DIR${NC}" >&2

    # 确保结果目录存在
    if [ ! -d "$FIXED_RESULT_DIR" ]; then
        echo -e "${YELLOW}[!] 结果目录不存在，尝试创建...${NC}" >&2
        mkdir -p "$FIXED_RESULT_DIR"
        chmod -R 777 "$FIXED_RESULT_DIR"
    fi

    # 检查结果目录是否可写
    if [ ! -w "$FIXED_RESULT_DIR" ]; then
        echo -e "${RED}[✗] 结果目录 $FIXED_RESULT_DIR 不可写！${NC}" >&2
        echo -e "${YELLOW}[!] 尝试修复权限...${NC}" >&2
        chmod -R 777 "$FIXED_RESULT_DIR"
        
        if [ ! -w "$FIXED_RESULT_DIR" ]; then
            echo -e "${RED}[✗] 权限修复失败，请手动执行: sudo chmod -R 777 $FIXED_RESULT_DIR${NC}" >&2
            exit 1
        fi
    fi

    echo -e "${GREEN}[✓] 结果目录准备就绪${NC}" >&2
}

# 生成带时间戳的结果文件名
generate_result_filename() {
    local scan_type="$1"
    local timestamp=$(date +"%Y%m%d_%H%M%S")
    echo "$FIXED_RESULT_DIR/${scan_type}_scan_${timestamp}.txt"
}

# 检查 Trivy 缓存目录和结果目录
check_trivy_env() {
    echo -e "${BLUE}[*] 正在检查 Trivy 运行环境...${NC}" >&2
    # 初始化缓存目录
    CACHE_DIR=$(get_cache_dir)
    DB_DIR="$CACHE_DIR/db"
    # 初始化结果目录
    create_results_dir
    RESULTS_DIR="$FIXED_RESULT_DIR"
    # 检查数据库目录
    if [ -d "$DB_DIR" ] && [ -f "$DB_DIR/metadata.json" ]; then
        echo -e "${YELLOW}[!] 检测到本地Trivy缓存数据库。${NC}"
        show_env_info
        read -p "是否使用本地缓存离线检测？(回车=离线，n=在线升级，q=退出): " use_offline
        if [[ "$use_offline" == "q" || "$use_offline" == "Q" ]]; then
            echo -e "${GREEN}已退出。${NC}"; exit 0
        fi
        if [[ -z "$use_offline" || "$use_offline" == "y" || "$use_offline" == "Y" ]]; then
            export TRIVY_OFFLINE_MODE=1
            echo -e "${GREEN}[✓] 已选择离线检测模式，将跳过数据库更新。${NC}"
        else
            export TRIVY_OFFLINE_MODE=0
            echo -e "${GREEN}[✓] 已选择在线升级模式。${NC}"
            read -p "是否现在自动更新数据库？(y=自动更新/n=返回菜单/q=退出): " auto_update
            if [[ "$auto_update" == "q" || "$auto_update" == "Q" ]]; then
                echo -e "${GREEN}已退出。${NC}"; exit 0
            fi
            if [[ "$auto_update" == "y" || "$auto_update" == "Y" ]]; then
                update_database
                read -p "按回车返回菜单..." tmp
                if [[ "$tmp" == "q" || "$tmp" == "Q" ]]; then echo -e "${GREEN}已退出。${NC}"; exit 0; fi
            fi
        fi
    else
        export TRIVY_OFFLINE_MODE=0
        show_env_info
    fi
    echo -e "${GREEN}[✓] 运行环境检查完成${NC}" >&2
    echo >&2
    # 导出缓存目录变量供Trivy使用
    export TRIVY_CACHE_DIR="$CACHE_DIR"
}

# 显示欢迎信息
show_welcome() {
    clear
    echo -e "${GREEN}==== Trivy 安全扫描工具 ====${NC}"
    echo -e "${YELLOW}固定缓存目录: $CACHE_DIR${NC}"
    echo -e "${YELLOW}固定结果目录: $RESULTS_DIR${NC}"
    echo
    show_menu
}

# 显示菜单
show_menu() {
    echo "请选择要执行的操作:"
    echo "1. 扫描本机漏洞（推荐，速度快）"
    echo "2. 扫描整机（全盘，最慢）"
    echo "3. 扫描指定目录"
    echo "4. 扫描常用服务漏洞"
    echo "5. 显示环境信息（缓存/结果目录）"
    echo "6. 修复目录权限（缓存+结果）"
    echo "7. 查看最近扫描结果"
    echo "8. 一键清理所有扫描结果文件"
    echo "9. 退出"
    echo
}

# 更新数据库
update_database() {
    echo -e "${GREEN}[+] 开始更新 Trivy 漏洞数据库...${NC}"
    echo
    if [ "$TRIVY_OFFLINE_MODE" = "1" ]; then
        echo -e "${YELLOW}[!] 离线模式，跳过数据库更新。${NC}"
        return 0
    fi
    # 优先尝试 trivy db update
    if trivy db update --cache-dir "$CACHE_DIR"; then
        echo
        echo -e "${GREEN}[✓] 数据库更新完成！${NC}"
    # 兼容老版本 trivy database update
    elif trivy database update --cache-dir "$CACHE_DIR"; then
        echo
        echo -e "${GREEN}[✓] 数据库更新完成！${NC}"
    # 极老版本 trivy --download-db-only
    elif trivy --download-db-only --cache-dir "$CACHE_DIR"; then
        echo
        echo -e "${GREEN}[✓] 数据库更新完成！${NC}"
    else
        echo
        echo -e "${YELLOW}当前 Trivy 版本不支持手动数据库更新，自动触发一次扫描以升级数据库...${NC}"
        # 自动触发一次fs扫描，仅用于升级数据库，不输出结果
        trivy fs --skip-db-update=false --cache-dir "$CACHE_DIR" /tmp >/dev/null 2>&1 || true
        echo -e "${GREEN}[✓] 已自动联网升级数据库，可直接进行后续扫描。${NC}"
    fi
    if [ -d "$DB_DIR" ]; then
        DB_VERSION=$(cat "$DB_DIR/metadata.json" 2>/dev/null | grep '"version":' | awk -F': ' '{print $2}' | tr -d ',')
        echo -e "${GREEN}[✓] 数据库版本: ${DB_VERSION:-未知}${NC}"
    else
        echo -e "${RED}[✗] 数据库目录不存在！${NC}"
        return 1
    fi
    return 0
}

# 扫描本机漏洞（仅常见系统目录，速度快）
scan_host() {
    echo -e "${GREEN}[+] 开始扫描本机常见系统目录（主机漏洞）...${NC}"
    echo -e "${YELLOW}[!] 仅检测常见系统目录，速度较快。${NC}"
    echo -e "${BLUE}[*] 扫描中请耐心等待，文件较多时可能需数分钟，按 Ctrl+C 可随时中断${NC}"
    echo
    RESULT_FILE=$(generate_result_filename "host")
    echo -e "${BLUE}[*] 结果将保存到: ${RESULT_FILE}${NC}"
    echo
    local skip_db=""
    if [ "$TRIVY_OFFLINE_MODE" = "1" ]; then
        skip_db="--skip-db-update"
    fi
    if [[ "$OSTYPE" == "darwin"* ]]; then
        TARGETS=(/System /usr /bin /sbin /etc)
    else
        TARGETS=(/etc /usr /bin /sbin)
    fi
    SCAN_TARGETS=()
    for d in "${TARGETS[@]}"; do
        if [ -d "$d" ] && [ "$(ls -A $d 2>/dev/null | wc -l)" -gt 0 ]; then
            SCAN_TARGETS+=("$d")
        fi
    done
    if [ ${#SCAN_TARGETS[@]} -eq 0 ]; then
        echo -e "${YELLOW}[!] 未检测到可用的系统目录用于扫描。${NC}"
        return 1
    fi
    echo -e "${BLUE}[*] 实际扫描目录: ${SCAN_TARGETS[*]}${NC}"
    echo
    # 捕获trivy stderr到临时文件
    local TRIVY_ERR=$(mktemp)
    if trivy fs \
        --scanners vuln \
        --severity CRITICAL,HIGH \
        --ignore-unfixed \
        $skip_db \
        --progress \
        --cache-dir "$CACHE_DIR" \
        "${SCAN_TARGETS[@]}" 2> "$TRIVY_ERR" | tee "$RESULT_FILE"; then
        if [ -s "$RESULT_FILE" ]; then
            # 检查扫描结果是否全为'-'（无有效扫描对象）
            if grep -qE '\|\s*-\s*\|\s*-\s*\|\s*-\s*\|' "$RESULT_FILE"; then
                echo -e "${YELLOW}[!] 未发现可识别的漏洞对象，建议扫描系统目录或包含依赖清单的应用目录。${NC}"
            else
                echo -e "${GREEN}[✓] 扫描完成，结果已保存到: ${RESULT_FILE}${NC}"
                echo -e "${BLUE}[*] 可将此文件下载后提供给AI分析漏洞，或通过‘查看最近扫描结果’菜单查看。${NC}"
            fi
            read -p "按回车返回主菜单..." _
        else
            echo -e "${YELLOW}[!] 扫描完成，但未发现漏洞或扫描被中断。${NC}"
            echo -e "${YELLOW}如需进一步分析，请使用‘查看最近扫描结果’菜单。${NC}"
            [ -f "$RESULT_FILE" ] && rm "$RESULT_FILE"
        fi
    else
        local code=$?
        echo -e "${RED}[✗] 扫描失败！Trivy 返回码: $code${NC}"
        echo -e "${YELLOW}[!] Trivy 错误日志如下:${NC}"
        cat "$TRIVY_ERR"
        echo -e "${YELLOW}如遇网络或权限问题，请检查环境或尝试修复目录权限。${NC}"
        [ -f "$RESULT_FILE" ] && rm "$RESULT_FILE"
        rm -f "$TRIVY_ERR"
        return 1
    fi
    rm -f "$TRIVY_ERR"
    return 0
}

# 扫描根文件系统（带权限跳过）
scan_rootfs() {
    echo -e "${GREEN}[+] 开始扫描根文件系统...${NC}"
    echo -e "${YELLOW}[!] 扫描可能需要几分钟时间，请耐心等待...${NC}"
    echo
    RESULT_FILE=$(generate_result_filename "rootfs")
    echo -e "${BLUE}[*] 结果将保存到: ${RESULT_FILE}${NC}"
    echo
    local skip_db=""
    if [ "$TRIVY_OFFLINE_MODE" = "1" ]; then
        skip_db="--skip-db-update"
    fi
    if trivy fs \
        --scanners vuln \
        --severity CRITICAL,HIGH \
        --ignore-unfixed \
        $skip_db \
        --cache-dir "$CACHE_DIR" \
        --skip-files "/www/wwwroot/*" \
        --skip-files "/sys/*" \
        --skip-files "/proc/*" \
        --skip-files "/dev/*" \
        --skip-files "/run/*" \
        --skip-files "/tmp/*" \
        / | tee "$RESULT_FILE"; then
        if [ -s "$RESULT_FILE" ]; then
            if grep -qE '\|\s*-\s*\|\s*-\s*\|\s*-\s*\|' "$RESULT_FILE"; then
                echo -e "${YELLOW}[!] 未发现可识别的漏洞对象，建议扫描系统目录或包含依赖清单的应用目录。${NC}"
            else
                echo -e "${GREEN}[✓] 扫描完成，结果已保存到: ${RESULT_FILE}${NC}"
                echo -e "${BLUE}[*] 可将此文件下载后提供给AI分析漏洞，或通过‘查看最近扫描结果’菜单查看。${NC}"
            fi
            read -p "按回车返回主菜单..." _
        else
            echo -e "${YELLOW}[!] 扫描完成，但未发现漏洞或扫描被中断。${NC}"
            echo -e "${YELLOW}如需进一步分析，请使用‘查看最近扫描结果’菜单。${NC}"
            [ -f "$RESULT_FILE" ] && rm "$RESULT_FILE"
        fi
    else
        echo -e "${RED}[✗] 扫描失败！${NC}"
        echo -e "${YELLOW}如遇网络或权限问题，请检查环境或尝试修复目录权限。${NC}"
        [ -f "$RESULT_FILE" ] && rm "$RESULT_FILE"
        return 1
    fi
    return 0
}

# 扫描指定目录
scan_directory() {
    echo -e "${BLUE}[*] 请输入要扫描的目录 (默认为 /，q=退出):${NC}"
    read -p "目录: " SCAN_DIR
    if [ "$SCAN_DIR" = "q" ] || [ "$SCAN_DIR" = "Q" ]; then
        echo -e "${GREEN}已退出。${NC}"; return 0
    fi
    if [ -z "$SCAN_DIR" ]; then
        SCAN_DIR="/"
    fi
    if [ ! -d "$SCAN_DIR" ]; then
        echo -e "${RED}[✗] 目录不存在: $SCAN_DIR${NC}"
        return 1
    fi
    echo -e "${GREEN}[+] 开始扫描目录: $SCAN_DIR${NC}"
    echo -e "${YELLOW}[!] 扫描可能需要几分钟时间，请耐心等待...${NC}"
    echo
    RESULT_FILE=$(generate_result_filename "directory")
    echo -e "${BLUE}[*] 结果将保存到: ${RESULT_FILE}${NC}"
    echo
    local skip_db=""
    if [ "$TRIVY_OFFLINE_MODE" = "1" ]; then
        skip_db="--skip-db-update"
    fi
    if trivy fs \
        --scanners vuln \
        --severity CRITICAL,HIGH \
        --ignore-unfixed \
        $skip_db \
        --cache-dir "$CACHE_DIR" \
        "$SCAN_DIR" | tee "$RESULT_FILE"; then
        if [ -s "$RESULT_FILE" ]; then
            if grep -qE '\|\s*-\s*\|\s*-\s*\|\s*-\s*\|' "$RESULT_FILE"; then
                echo -e "${YELLOW}[!] 未发现可识别的漏洞对象，建议扫描系统目录或包含依赖清单的应用目录。${NC}"
            else
                echo -e "${GREEN}[✓] 扫描完成，结果已保存到: ${RESULT_FILE}${NC}"
                echo -e "${BLUE}[*] 可将此文件下载后提供给AI分析漏洞，或通过‘查看最近扫描结果’菜单查看。${NC}"
            fi
            read -p "按回车返回主菜单..." _
        else
            echo -e "${YELLOW}[!] 扫描完成，但未发现漏洞或扫描被中断。${NC}"
            echo -e "${YELLOW}如需进一步分析，请使用‘查看最近扫描结果’菜单。${NC}"
            [ -f "$RESULT_FILE" ] && rm "$RESULT_FILE"
        fi
    else
        echo -e "${RED}[✗] 扫描失败！${NC}"
        echo -e "${YELLOW}如遇网络或权限问题，请检查环境或尝试修复目录权限。${NC}"
        [ -f "$RESULT_FILE" ] && rm "$RESULT_FILE"
        return 1
    fi
    return 0
}

# 扫描常用服务漏洞（自动检测常见服务路径）
scan_services() {
    echo -e "${GREEN}[+] 开始扫描常用服务漏洞...${NC}"
    echo -e "${YELLOW}[!] 自动检测并扫描常见服务（nginx、mysql、redis、docker等）安装路径。${NC}"
    echo
    local skip_db=""
    if [ "$TRIVY_OFFLINE_MODE" = "1" ]; then
        skip_db="--skip-db-update"
    fi
    # 常见服务及其典型路径
    declare -A services
    services=(
        [nginx]="/etc/nginx /usr/sbin/nginx /usr/local/nginx"
        [mysql]="/etc/mysql /usr/sbin/mysqld /usr/local/mysql"
        [redis]="/etc/redis /usr/bin/redis-server /usr/local/bin/redis-server"
        [docker]="/etc/docker /usr/bin/dockerd /usr/local/bin/dockerd"
        [httpd]="/etc/httpd /usr/sbin/httpd /usr/local/apache2"
        [postgres]="/etc/postgresql /usr/bin/postgres /usr/local/pgsql"
        [mongodb]="/etc/mongod.conf /usr/bin/mongod /usr/local/bin/mongod"
    )
    local found=0
    for svc in "${!services[@]}"; do
        for path in ${services[$svc]}; do
            if [ -e "$path" ]; then
                found=1
                RESULT_FILE=$(generate_result_filename "$svc")
                echo -e "${BLUE}[*] 检测到 $svc 路径: $path${NC}"
                echo -e "${BLUE}[*] 结果将保存到: ${RESULT_FILE}${NC}"
                if [ -d "$path" ] || [[ "$path" == *.conf ]]; then
                    trivy fs \
                        --scanners vuln \
                        --severity CRITICAL,HIGH \
                        --ignore-unfixed \
                        $skip_db \
                        --cache-dir "$CACHE_DIR" \
                        "$path" | tee "$RESULT_FILE"
                else
                    trivy fs \
                        --scanners vuln \
                        --severity CRITICAL,HIGH \
                        --ignore-unfixed \
                        $skip_db \
                        --cache-dir "$CACHE_DIR" \
                        "$(dirname $path)" | tee "$RESULT_FILE"
                fi
                if [ -s "$RESULT_FILE" ]; then
                    # 检查扫描结果是否全为'-'（无有效扫描对象）
                    if grep -qE '\|\s*-\s*\|\s*-\s*\|\s*-\s*\|' "$RESULT_FILE"; then
                        echo -e "${YELLOW}[!] 未发现可识别的漏洞对象，建议扫描系统目录或包含依赖清单的应用目录。${NC}"
                    else
                        echo -e "${GREEN}[✓] $svc 扫描完成，结果已保存到: ${RESULT_FILE}${NC}"
                        echo -e "${BLUE}[*] 可将此文件下载后提供给AI分析漏洞，或通过‘查看最近扫描结果’菜单查看。${NC}"
                    fi
                    read -p "按回车返回主菜单..." _
                else
                    echo -e "${YELLOW}[!] $svc 扫描完成，但未发现漏洞或扫描被中断。${NC}"
                    echo -e "${YELLOW}如需进一步分析，请使用‘查看最近扫描结果’菜单。${NC}"
                    [ -f "$RESULT_FILE" ] && rm "$RESULT_FILE"
                fi
                echo
            fi
        done
    done
    if [ $found -eq 0 ]; then
        echo -e "${YELLOW}[!] 未检测到常用服务安装路径，可用‘扫描指定目录’手动输入。${NC}"
    fi
    return 0
}

# 兼容 stat 和 date
get_perm() {
    local path="$1"
    if [[ "$OSTYPE" == "darwin"* ]]; then
        stat -f "%Sp" "$path"
    else
        stat -c "%A" "$path"
    fi
}
get_owner() {
    local path="$1"
    if [[ "$OSTYPE" == "darwin"* ]]; then
        stat -f "%Su:%Sg" "$path"
    else
        stat -c "%U:%G" "$path"
    fi
}
get_last_update() {
    local ts="$1"
    if [[ -z "$ts" ]]; then
        echo "未知"
        return
    fi
    if [[ "$OSTYPE" == "darwin"* ]]; then
        date -r "$ts" "+%Y-%m-%d %H:%M:%S" 2>/dev/null || echo "未知"
    else
        date -d "@$ts" "+%Y-%m-%d %H:%M:%S" 2>/dev/null || echo "未知"
    fi
}

# 显示环境信息（缓存/结果目录）
show_env_info() {
    echo -e "${BLUE}[*] 系统环境信息:${NC}"
    echo "缓存目录: $CACHE_DIR"
    echo "结果目录: $RESULTS_DIR"
    if [ -d "$CACHE_DIR" ]; then
        size=$(du -sh $CACHE_DIR 2>/dev/null | awk '{print $1}')
        perm=$(get_perm $CACHE_DIR)
        owner=$(get_owner $CACHE_DIR)
        echo -n "缓存目录详情: 大小: $size, 权限: $perm, 所有者: $owner"
        if [ -d "$DB_DIR" ]; then
            local ts=$(cat "$DB_DIR/metadata.json" 2>/dev/null | grep '"updatedAt":' | awk -F': ' '{print $2}' | tr -d ',"')
            echo -n ", 上次更新: $(get_last_update $ts)"
        fi
        echo
    else
        echo -e "${RED}缓存目录不存在！${NC}"
    fi
    if [ -d "$RESULTS_DIR" ]; then
        size=$(du -sh $RESULTS_DIR 2>/dev/null | awk '{print $1}')
        perm=$(get_perm $RESULTS_DIR)
        owner=$(get_owner $RESULTS_DIR)
        echo -n "结果目录详情: 大小: $size, 权限: $perm, 所有者: $owner"
        files=$(ls -t $RESULTS_DIR 2>/dev/null | head -n 5 | tr '\n' ' ')
        if [ -n "$files" ]; then
            echo -n ", 最近结果文件: $files"
        fi
        echo
    else
        echo -e "${RED}结果目录不存在！${NC}"
    fi
}

# 查看最近扫描结果
show_recent_results() {
    echo -e "${BLUE}[*] 最近的扫描结果:${NC}"
    
    if [ ! -d "$RESULTS_DIR" ]; then
        echo -e "${YELLOW}[!] 结果目录不存在，没有扫描结果${NC}"
        return
    fi
    
    # 列出最近5个结果文件
    recent_files=$(ls -t $RESULTS_DIR 2>/dev/null | head -n 5)
    
    if [ -z "$recent_files" ]; then
        echo -e "${YELLOW}[!] 没有找到扫描结果${NC}"
        return
    fi
    
    echo "$recent_files" | awk '{print NR ". " $0}'
    echo
    
    read -p "输入编号查看详情 (0返回，q=退出): " choice
    if [ "$choice" = "q" ] || [ "$choice" = "Q" ]; then
        echo -e "${GREEN}已退出。${NC}"; return
    fi
    if [ "$choice" = "0" ]; then
        return
    fi
    
    # 获取选择的文件
    selected_file=$(echo "$recent_files" | sed -n "${choice}p")
    
    if [ -z "$selected_file" ]; then
        echo -e "${RED}[!] 无效选择${NC}"
        return
    fi
    
    selected_path="$RESULTS_DIR/$selected_file"
    
    if [ -f "$selected_path" ]; then
        echo -e "${BLUE}[*] 文件内容: ${selected_file}${NC}"
        echo "----------------------------------------"
        cat "$selected_path"
        echo "----------------------------------------"
        echo -e "${BLUE}[*] 文件路径: ${selected_path}${NC}"
        echo -e "${BLUE}[*] 可通过此路径下载文件供AI分析${NC}"
    else
        echo -e "${RED}[!] 文件不存在${NC}"
    fi
}

# 修复目录权限（缓存+结果）
fix_permissions() {
    echo -e "${BLUE}[*] 正在修复所有目录权限...${NC}"
    
    # 修复缓存目录
    echo -e "\n${YELLOW}[*] 处理缓存目录: $CACHE_DIR${NC}"
    mkdir -p "$CACHE_DIR"
    chmod -R 777 "$CACHE_DIR"
    echo -e "${GREEN}[✓] 缓存目录权限已修复${NC}"
    
    # 修复结果目录
    echo -e "\n${YELLOW}[*] 处理结果目录: $RESULTS_DIR${NC}"
    mkdir -p "$RESULTS_DIR"
    chmod -R 777 "$RESULTS_DIR"
    echo -e "${GREEN}[✓] 结果目录权限已修复${NC}"
    
    echo -e "\n${GREEN}[✓] 所有目录权限修复完成${NC}"
}

# 检查trivy是否安装，未安装则提示用户先运行安装脚本
check_trivy_installed() {
    if ! command -v trivy &> /dev/null; then
        echo -e "${RED}[✗] 未检测到Trivy命令！${NC}"
        local script_dir="$(cd "$(dirname "$0")" && pwd)"
        if [ -f "$script_dir/install_trivy.sh" ]; then
            read -p "未检测到Trivy，是否现在自动执行 install_trivy.sh 进行安装？(y/n): " auto_install
            if [[ "$auto_install" == "y" || "$auto_install" == "Y" ]]; then
                bash "$script_dir/install_trivy.sh"
                echo -e "${GREEN}[✓] Trivy 安装脚本已执行，请重新运行本脚本。${NC}"
                exit 0
            else
                echo -e "${YELLOW}未安装Trivy，程序退出。${NC}"
                exit 1
            fi
        else
            echo -e "${RED}未找到 install_trivy.sh，无法自动安装Trivy。${NC}"
            exit 1
        fi
    fi
}

# 一键清理所有扫描结果文件
clean_results() {
    echo -e "${YELLOW}警告：此操作将删除结果目录（$RESULTS_DIR）下的所有扫描结果文件！${NC}"
    read -p "确定要清空所有扫描结果吗？(y/n): " confirm
    if [[ "$confirm" == "y" || "$confirm" == "Y" ]]; then
        if [ -d "$RESULTS_DIR" ]; then
            rm -rf "$RESULTS_DIR"/*
            echo -e "${GREEN}[✓] 所有扫描结果已清空。${NC}"
        else
            echo -e "${YELLOW}[!] 结果目录不存在，无需清理。${NC}"
        fi
    else
        echo -e "${YELLOW}已取消清理操作。${NC}"
    fi
    read -p "按回车返回主菜单..." _
}

# 主程序入口
main() {
    check_trivy_installed
    check_command trivy
    check_command du
    check_command date
    check_command stat
    check_trivy_env
    CACHE_DIR=$(get_cache_dir)
    DB_DIR="$CACHE_DIR/db"
    RESULTS_DIR="$FIXED_RESULT_DIR"
    export CACHE_DIR
    export RESULTS_DIR
    while true; do
        show_welcome
        read -p "请输入选项编号: " opt
        case "$opt" in
            1) scan_host ;;
            2) scan_rootfs ;;
            3) scan_directory ;;
            4) scan_services ;;
            5) show_env_info; read -p "按回车返回菜单..." ;;
            6) fix_permissions ;;
            7) show_recent_results; read -p "按回车返回菜单..." ;;
            8) clean_results ;;
            9) echo -e "${GREEN}退出脚本。${NC}"; exit 0 ;;
            *) echo -e "${RED}无效选项，请重新输入。${NC}" ;;
        esac
    done
}

# 启动主程序
main
